import numpy as np
import cv2

import caffe


MODEL_FILE = 'mbv2_yolov3_face_conv.prototxt'
CAFFEMODEL_FILE = 'mbv2_yolov3_face_conv.caffemodel'
DARKNET_WEIGHTS_FILE = 'mbv2_yolov3_face_final.weights'



if __name__ =='__main__':

    # Read weights from file (dtype=int32)
    netWeightsInt = np.fromfile(DARKNET_WEIGHTS_FILE, dtype=np.int32)
    # First 4 entries are major, minor, revision, and seen(uint64_t) with int type.
    # int transpose = (major > 1000) || (minor > 1000);
    transpose = (netWeightsInt[0]>1000 or netWeightsInt[1]>1000)
    print('transpose: ', transpose)

    # Read weights from file (dtype=float32)
    netWeightsFloat = np.fromfile(DARKNET_WEIGHTS_FILE, dtype=np.float32)
    netWeights = netWeightsFloat[5:]
    print(netWeights.shape)


    # Init Caffe Network
    net = caffe.Net(MODEL_FILE, caffe.TEST)

    # Params
    params_keys = net.params.keys()
    params_vals = net.params.values()


    # Weights
    count = 0
    for param in params_keys:

        lidx = list(net._layer_names).index(param)
        layer = net.layers[lidx]

        if count == netWeights.shape[0] and (layer.type != 'BatchNorm' and layer.type != 'Scale'):
            print('WARNING: No Weights Left for Remaining Layers!')

        if layer.type == 'Convolution':

            # bias
            if len(net.params[param]) > 1:
                bias_dims = net.params[param][1].data.shape
            else:
                bias_dims = (net.params[param][0].shape[0], )
            bias_lens = np.prod(bias_dims)
            print('conv bias: ', bias_dims, bias_lens)

            conv_bias = np.reshape(netWeights[count:count+bias_lens], bias_dims)
            count = count + bias_lens

            if len(net.params[param]) > 1:  # Convolution with bias_term = True.
                assert(bias_dims == net.params[param][1].data.shape)
                net.params[param][1].data[...] = conv_bias
                conv_bias = None


            # If BatchNorm Layer follows
            if lidx+1 < len(net.layers) and net.layers[lidx+1].type == 'BatchNorm':
                bn_dims = (3, net.params[param][0].data.shape[0])
                bn_lens = np.prod(bn_dims)
                print('batch norm: ', bn_dims, bn_lens)

                batchnorm = np.reshape(netWeights[count:count+bn_lens], bn_dims)
                count = count + bn_lens


            # Conv Weights
            dims = net.params[param][0].data.shape
            lens = np.prod(dims)
            print('conv wgts: ', dims, lens)

            net.params[param][0].data[...] = np.reshape(netWeights[count:count+lens], dims)
            count = count + lens


        elif layer.type == 'BatchNorm':

            print('BatchNorm')
            net.params[param][0].data[...] = batchnorm[1]	# mean
            net.params[param][1].data[...] = batchnorm[2]	# variance
            net.params[param][2].data[...] = 1.0		# scale factor


        elif layer.type == 'Scale':

            print('Scale')
            net.params[param][0].data[...] = batchnorm[0]       # scale
            batchnorm = None

            if len(net.params[param]) > 1:  # Move convolution bias into Scale Layer
                print('conv bias in scale')
                net.params[param][1].data[...] = conv_bias
                conv_bias = None


        elif layer.type == 'InnerProduct':

            print('InnerProduct')

            # bias
            bias_dims = net.params[param][1].data.shape
            bias_lens = np.prod(bias_dims)
            net.params[param][1].data[...] = np.reshape(netWeights[count:count+bias_lens], bias_dims)
            count = count + bias_lens

            # weights
            dims = net.params[param][0].data.shape
            lens = np.prod(dims)

            if transpose:
                net.params[param][0].data[...] = np.reshape(netWeights[count:count+lens], (dims[1],dims[0])).transpose()
            else:
                net.params[param][0].data[...] = np.reshape(netWeights[count:count+lens], dims)
            count = count + lens


        else:
            print('WARNING: unsupported layer -- {}'.format(param))


    print('original weights shape: ', netWeights.shape, 'used weights: ', count)
    net.save(CAFFEMODEL_FILE)

